import random
import json
from locust import task, TaskSet, HttpUser, between, FastHttpUser, events
from locust.env import Environment
from gevent._semaphore import Semaphore

all_locusts_spawned = Semaphore()
all_locusts_spawned.acquire()  # 阻塞线程

def on_hatch_complete(user_count, **kwargs):
    """
    Select_task类的钩子方法
    """
    print(f"{user_count}个虚拟用户产生完成")
    # 创建钩子方法
    all_locusts_spawned.release()
events.spawning_complete.add_listener(on_hatch_complete)

class subEnvironment(Environment):
    def __init__(self):
        super().__init__()
        self.base1_url = None  # 你可以在这里添加更多的环境变量
        self.base2_url = None

my_env = subEnvironment()
n = 0


class query_order(TaskSet):
    weight = 1
    host = "https://tmstest.gree.com"

    def getRandomString(self, length, num=True):
        '获取一个随机字符串，length：随机结果字符串长度；num：随机结果是否包含数字(默认包含)，True：包含，False：不包含；'
        strList = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's',
                   't', 'u', 'v', 'w', 'x', 'y', 'z']
        numList = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        result = ''
        for i in range(length):
            if num:
                result += str(random.choice(strList + numList))
            else:
                result += str(random.choice(strList))
        return result

    @task
    def createOrder(self):
        order = self.getRandomString(18)

        param = {
            "uid": "4056",
            "tmsOrderCodes": [
                order
            ],
            "pageIndex": 1,
            "pageSize": 10,
            "tmsOrderTimeBegin": "2024-08-06 00:00:00",
            "tmsOrderTimeEnd": "2024-11-06 23:59:59"
        }
        with self.client.post("/tms/business/order/work", data=json.dumps(param), headers=headers,
                              catch_response=True) as res:
            pass
            # if res.status_code != 200:
            #     print("oms-data_order入参：%s, 返回值:%s" % (param, res.text))

    def on_start(self):
        global n
        n += 1
        print("%s个虚拟用户开始启动" % n)
        all_locusts_spawned.wait()  # 同步锁等待
class Tms(FastHttpUser):
    global headers
    headers = {
        "Content-Type": "application/json;charset=UTF-8",
        "uid":4056,
        "authorization":"256e2b7b-da54-46ab-b348-0b2d00016984",
        "systemflag":10


    }
    host = ""
    wait_time = between(0, 0)


    # class query_order_detail(TaskSet):
    #     weight = 3
    #     my_env.base_url2 = "http://172.28.25.82:39615"
    #     def getRandomString(self, length, num=True):
    #         '获取一个随机字符串，length：随机结果字符串长度；num：随机结果是否包含数字(默认包含)，True：包含，False：不包含；'
    #         strList = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's',
    #                    't', 'u', 'v', 'w', 'x', 'y', 'z']
    #         numList = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    #         result = ''
    #         for i in range(length):
    #             if num:
    #                 result += str(random.choice(strList + numList))
    #             else:
    #                 result += str(random.choice(strList))
    #         return result
    #
    #     def on_start(self):
    #         global n
    #         n += 1
    #         print("%s个虚拟用户开始启动" % n)
    #         all_locusts_spawned.wait()  # 同步锁等待
    #
    #
    #     @task
    #     def createOrder(self):
    #         order = self.getRandomString(18)
    #
    #         param = {
    #                 "uid": "4056",
    #                 "warehouseCodeList": [],
    #                 "carrierCodeList": [],
    #                 "platformCodeList": [],
    #                 "materialCode": order,
    #                 "pageIndex": 1,
    #                 "pageSize": 10
    #             }
    #         with self.client.post(f"{my_env.base_url2}/tms/business/product_config/lis", data=json.dumps(param), headers=headers,
    #                               catch_response=True) as res:
    #             if res.status_code != 200:
    #                 print("oms-data_order入参：%s, 返回值:%s" % (param, res.text))

    tasks = [query_order]
    # tasks = [query_order, query_order_detail]


if __name__ == '__main__':
    import subprocess
    subprocess.run(
        ["locust", "-f", "test.py"])